package com.pocket.insight.base.context;

import com.pocket.insight.base.exception.SysException;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.BooleanUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.Connection;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Function;

/**
 * 描述：
 *
 * @author shuize
 * @create 2024/11/29
 */
public class SQLExecContext {
    private static final Logger LOGGER = LoggerFactory.getLogger(SQLExecContext.class);

    private static ThreadLocal<List<Object>>       connectionsThreadLocal    = new ThreadLocal<>();
    private static ConcurrentHashMap<Long, Object> jobThreadCurrentStatement = new ConcurrentHashMap<>();
    private static ThreadLocal<Map<Long, Object>>  connToVendor              = new ThreadLocal<>();
    private static ThreadLocal<List<Runnable>>     rollbackHandlerList       = new ThreadLocal<>();
    private static ThreadLocal<Boolean> etl = new ThreadLocal<>();

    public static void setEtl(Boolean isEtl) {
        etl.set(isEtl);
    }

    public static boolean isEtl() {
        return BooleanUtils.isTrue(etl.get());
    }

    public static void commitAll() {
        CollectionUtils.emptyIfNull(connectionsThreadLocal.get()).forEach(one -> {
            try {
                if (!((Connection)one).getAutoCommit()) {
                    ((Connection)one).commit();
                }
            } catch (Exception e) {
                LOGGER.error(e.getMessage(), e);
                throw new SysException(String.format("ERROR_FAIL_TO_WRITE_TO_DATABASE: ", e.getMessage()));
            }
        });
    }

    public static void rollbackAll() {
        CollectionUtils.emptyIfNull(connectionsThreadLocal.get()).forEach(one -> {
            try {
                if (!((Connection)one).getAutoCommit()) {
                    ((Connection)one).rollback();
                }
            } catch (Exception e) {
                LOGGER.error(e.getMessage(), e);
                throw new SysException(String.format("ERROR_FAIL_TO_WRITE_TO_DATABASE: ", e.getMessage()));
            }
        });
        CollectionUtils.emptyIfNull(rollbackHandlerList.get()).forEach(runnable -> {
            if (runnable != null) {
                runnable.run();
            }
        });
    }

    public static void addConnection(Connection connection) {
        if (connectionsThreadLocal.get() == null) {
            connectionsThreadLocal.set(new ArrayList());
        }
        (connectionsThreadLocal.get()).add(connection);
    }

    public static void removeConnection(Connection connection) {
        List<Object> conns = connectionsThreadLocal.get();
        if (conns != null) {
            conns.remove(connection);
        }
    }

    public static Object getVendor(Long connId, Function<Void, Object> vendorFactory) {
        if (connToVendor.get() == null) {
            connToVendor.set(new HashMap<>());
        }
        if (!((Map)connToVendor.get()).containsKey(connId)) {
            ((Map)connToVendor.get()).put(connId, vendorFactory.apply(null));
        }
        return ((Map)connToVendor.get()).get(connId);
    }

    public static void addRollbackHandler(Runnable rollbackHandler) {
        if (rollbackHandler == null) {
            return;
        }
        if (rollbackHandlerList.get() == null) {
            rollbackHandlerList.set(new ArrayList<>());
        }
        ((List<Runnable>)rollbackHandlerList.get()).add(rollbackHandler);
    }

    public static void setCurrentStatement(Statement statement) {
        jobThreadCurrentStatement.put(Long.valueOf(Thread.currentThread().getId()), statement);
    }

    public static void cancelCurrentStatement(Long threadId) {
        Object stm = jobThreadCurrentStatement.get(threadId);
        if (stm != null) {
            try {
                ((Statement)stm).cancel();
            } catch (Exception e) {
                LOGGER.error(e.getMessage(), e);
            }
        }
    }

    public static void remove() {
        CollectionUtils.emptyIfNull(connectionsThreadLocal.get()).forEach(one -> {
            try {
                ((Connection)one).close();
            } catch (Exception e) {
                LOGGER.error(e.getMessage(), e);
            }
        });
        connectionsThreadLocal.remove();
        jobThreadCurrentStatement.clear();
        connToVendor.remove();
        rollbackHandlerList.remove();
        etl.remove();
    }
}
